package com.kanq.service.impl;

import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import com.kanq.documentLoader.result.Section;
import com.kanq.mapper.KanqAiFileChunkMapper;
import com.kanq.pojo.entity.KanqAiFileChunk;
import com.kanq.pojo.entity.KanqAiFileChunkKeywords;
import com.kanq.service.KanqAiFileChunkKeywordsService;
import com.kanq.service.KanqAiFileChunkService;
import com.kanq.service.QdrantService;
import jakarta.annotation.Resource;
import org.springframework.stereotype.Service;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

@Service
public class KanqAiFileChunkServiceImpl extends ServiceImpl<KanqAiFileChunkMapper, KanqAiFileChunk> implements KanqAiFileChunkService {

    @Resource
    private KanqAiFileChunkKeywordsService kanqAiFileChunkKeywordsService;

    @Resource
    private QdrantService qdrantService;

    @Resource
    private KanqAiFileChunkMapper kanqAiFileChunkMapper;


    /**
     * 保存文件分块及其关键字
     *
     * @param fileId   文件ID
     * @param content  分块内容
     * @param keywords 关键字列表
     * @return 保存的分块ID
     */
    @Override
    public KanqAiFileChunk saveChunkAndKeywords(long fileId, String content, List<String> keywords) {
        KanqAiFileChunk chunk = new KanqAiFileChunk();
        chunk.setFileId(Long.valueOf(fileId));
        chunk.setContent(content);
        baseMapper.insert(chunk);
        List<KanqAiFileChunkKeywords> keywordsList = new ArrayList<>();
        for (String keyword : keywords) {
            if (keyword.trim().isEmpty()) {
                continue;
            }
            KanqAiFileChunkKeywords kbFileChunkKeywords = new KanqAiFileChunkKeywords();
            kbFileChunkKeywords.setChunkId(chunk.getId());
            kbFileChunkKeywords.setKeyword(keyword);
            keywordsList.add(kbFileChunkKeywords);
        }
        kanqAiFileChunkKeywordsService.saveKeywords(keywordsList);
        return chunk;
    }

    @Override
    public List<KanqAiFileChunk> fullTextSearch(List<String> keywords) {
        Map<String, String> params = new HashMap<>();
        StringBuilder keyword = new StringBuilder();
        for (int i = 0; i < keywords.size(); i++) {
            if (i == keywords.size() - 1) {
                keyword.append(keywords.get(i));
            } else {
                keyword.append(keywords.get(i)).append(" | ");
            }
        }
        params.put("keywords", keyword.toString());

        return kanqAiFileChunkMapper.searchByKeywords(params);
    }


    /**
     * 批量插入文件分块并生成向量
     *
     * @param fileId   文件ID
     * @param contents 分块内容列表
     * @return 插入是否成功
     */
    @Override
    public boolean saveFileChunk(long fileId, String fileName, List<Section> contents) {
        List<KanqAiFileChunk> chunks = new ArrayList<>();
        for (Section content : contents) {
            KanqAiFileChunk chunk = new KanqAiFileChunk();
            chunk.setFileId(Long.valueOf(fileId));
            chunk.setContent(content.toString());
            chunks.add(chunk);
        }

        this.saveBatch(chunks, 100);

        for (KanqAiFileChunk chunk : chunks) {
            try {
                qdrantService.insertVectorInString(chunk.getId(), chunk.getContent(), fileId, fileName);
            } catch (Exception e) {
                return false;
            }
        }
        return true;
    }

    @Override
    public List<KanqAiFileChunk> getSegDataPage(long fileId) {
        return baseMapper.getSegDataPage(fileId);
    }
}
